/*
	Copyright (C) 2011 Salil Bhagurkar

	This file is part of illusion

	illusion is free software: you can redistribute it and/or modify
	it under the terms of the GNU Lesser General Public License as
	published by the Free Software Foundation, either version 3 of
	the License, or (at your option) any later version.

	illusion is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Lesser General Public License for more details.

	You should have received a copy of the GNU Lesser General Public
	License along with illusion. 
	If not, see <http://www.gnu.org/licenses/>.
*/

#include <klib/lib.h>
#include <kernel/list.h>
#include <kernel/page.h>
#include <apimod/apimod.h>
#include <kernel/errors.h>
#include <kernel/build.h>
#include <kernel/config.h>


/*
 * A port of the kernel memory allocator for user space
 */

#define BIT_DIV (8 * sizeof(u32))


static int get_bit(void *bitmap, uint_t nr)
{
	uint_t i = nr / BIT_DIV;
	uint_t bit = nr % BIT_DIV;
	return (((u32 *)bitmap)[i] & (1UL << bit)) ? 1 : 0;
}

static void set_bit(void *bitmap, uint_t nr)
{
	uint_t i = nr / BIT_DIV;
	uint_t bit = nr % BIT_DIV;
	((u32 *)bitmap)[i] |= (1UL << bit);
}

static void clear_bit(void *bitmap, uint_t nr)
{
	uint_t i = nr / BIT_DIV;
	uint_t bit = nr % BIT_DIV;
	((u32 *)bitmap)[i] &= ~(1UL << bit);
}


#define PAGE_MASK 0xFFFFF000

//The sign at the beginning of each page
#define MEM_PAGE_SIGN 0x32313132
//The size of one block that can be allocated
#define MEM_BLOCK_SIZE 32
//The total number of blocks in one page
#define MEM_PAGE_BLOCKS (PAGE_SIZE / MEM_BLOCK_SIZE)
//The number of u32s that are required for the bitmap
#define MEM_PAGE_MAP_WORDS (MEM_PAGE_BLOCKS / (8 * sizeof(u32)))
//The number of blocks required to reserve space for the header itself
#define MEM_PAGE_HEADER_BLOCKS (sizeof(struct mem_page) / MEM_BLOCK_SIZE + !!((sizeof(struct mem_page)) % MEM_BLOCK_SIZE))
//The number of blocks that are really available for allocation
#define MEM_PAGE_AVL_BLOCKS (MEM_PAGE_BLOCKS - MEM_PAGE_HEADER_BLOCKS)
//The memory available in one page
#define MEM_PAGE_AVL_MEM (MEM_PAGE_AVL_BLOCKS * MEM_BLOCK_SIZE)


/***********Adds a signature to each allocation for corruption testing*******/
#ifdef CONFIG_MALLOC_GUARD_SIGN
static void alloc_enter(uint_t *new_size, uint_t size)
{
	*new_size = size + 2 * sizeof(u32);
}

static void alloc_exit(void **new_mem, void *mem, uint_t size)
{
	u32 *start_addr = (u32 *)mem;
	u32 *end_addr = mem + size - sizeof(u32);
	*start_addr = MEM_PAGE_SIGN;
	*end_addr = MEM_PAGE_SIGN;
	*new_mem = mem + sizeof(u32);
}

static void free_verify(void **real_mem, uint_t *real_size, void *mem, uint_t size)
{
	*real_size = size + 2 * sizeof(u32);
	u32 *start_addr = (u32 *)(mem - sizeof(u32));
	u32 *end_addr = (u32 *)(((void *)start_addr) + (*real_size) - sizeof(u32));
	if((*start_addr) != MEM_PAGE_SIGN || (*end_addr) != MEM_PAGE_SIGN) {
		//Throw an error with what the caller thinks are the right
		//attributes of the allocation
		insane(ESTINV, "Signature mismatch: %x, %x, [0x%x, %d B]\n",
				(*start_addr), (*end_addr), mem, size);
	}
	*real_mem = mem - sizeof(u32);
}
#else
#define alloc_enter
#define alloc_exit
#define free_verify
#endif
/****************************************************************************/


/*
 * The header of any memory page
 */
struct mem_page {
	u32 sign;
	u32 bitmap[MEM_PAGE_MAP_WORDS]; //Maps 128 32 byte blocks
	uint_t free_blocks; //Maximum free blocks in this page
	struct mem_page *next, *prev;
};


/*
 * The list of pages kept in pool when only a part of them is allocated
 */
static struct mem_page *pages = null, *pages_tail = null;

/*
 * Get the number of pages in the pages, pagesTail list
 * These pages are kept if they have partially free memory
 */
static uint_t pool_page_count()
{
	struct mem_page *page;
	uint_t nr_pages = 0;
	list_for(pages, page) {
		nr_pages++;
	}
	return nr_pages;
}


/*
 * Get the total number of free blocks summing up all
 * in each of the pages in the pool.
 */
static uint_t pool_free_blocks()
{
	struct mem_page *page;
	uint_t free_blocks = 0;
	list_for(pages, page) {
		free_blocks += page->free_blocks;
	}
	return free_blocks;
}

static uint_t malloc_pool_memory()
{
	return pool_page_count() * PAGE_SIZE;
}

static uint_t malloc_pool_free_memory()
{
	return pool_free_blocks() * MEM_BLOCK_SIZE;
}

/*
 * Initialize a page and add it to the list
 */
static void init_mem_page(struct mem_page *page)
{
	page->sign = MEM_PAGE_SIGN;
	memory_set(page->bitmap, 0, sizeof(page->bitmap));
	//Allocate the memory used by the header
	uint_t i;
	for(i = 0; i < MEM_PAGE_HEADER_BLOCKS; i++)
		set_bit(&page->bitmap, i);
	page->free_blocks = MEM_PAGE_AVL_BLOCKS;
	page->next = page->prev = null;
}


/*
 * Find a location where we could fit the required amount of free blocks
 */
static uint_t get_fit(void *bitmap, uint_t req_blocks)
{
	uint_t at = 0, i;
	uint_t len = 0;
	char in_contig = 0;
	for(i = 0; i < MEM_PAGE_BLOCKS; i++) {
		if(!get_bit(bitmap, i)) {
			if(!in_contig) {
				in_contig = 1;
				at = i;
			}
			len++;
			if(len >= req_blocks) {
				return at;
			}
		} else {
			in_contig = 0;
			len = 0;
		}
	}
	if(len >= req_blocks) {
		return at;
	}
	return 0;
}

/*
 * Find a page with enough free space from the list of pages
 */
static struct mem_page *find_page(uint_t size, uint_t *location)
{
	struct mem_page *page;
	uint_t fit;
	*location = 0;
	list_for(pages, page) {
		fit = get_fit(&page->bitmap, size);
		if(fit > 0) {
			*location = fit;
			return page;
		}
	}
	return null;
}

void *malloc(uint_t size)
{
	if(size == 0)
		return null;

	alloc_enter(&size, size);

	//If the size is greater than one page, allocate a page
	if(size > MEM_PAGE_AVL_MEM) {
		void *ret = get_pages((size / PAGE_SIZE) + !!(size % PAGE_SIZE));
		alloc_exit(&ret, ret, size);
		return ret;
	}
	uint_t size_blocks = size / MEM_BLOCK_SIZE + !!(size % MEM_BLOCK_SIZE);

	//Find if there is any memory available in the page pool
	uint_t index;
	struct mem_page *page = find_page(size_blocks, &index);

	if(!page) {
		//If none available, get a free page
		page = get_pages(1);
		init_mem_page(page);

		list_attach(pages, pages_tail, page);
		//Starting index of allocation will be after the header
		index = MEM_PAGE_HEADER_BLOCKS;
	}

	uint_t i;

	//Allocate the blocks
	uint_t end = index + size_blocks;
	for(i = index; i < end; i++) {
		set_bit(&page->bitmap, i);
	}

	page->free_blocks -= size_blocks;
	//If the page doesn't have any free memory left now, remove the page
	//It will come back once any memory from it is freed
	if(!(page->free_blocks)) {
		list_detach(pages, pages_tail, page);
	}

	void *ret = (void *)page + index * MEM_BLOCK_SIZE;

	alloc_exit(&ret, ret, size);

	return ret;
}

void free(void *addr, uint_t size)
{
	if(!addr || size == 0) {
		insane(EPARAM, "Null pointer or zero size passed to free: 0x%x of claimed size %d B\n", addr, size);
	}

	free_verify(&addr, &size, addr, size);

	if(size > MEM_PAGE_AVL_MEM) {
		uint_t nr_pages = size / PAGE_SIZE + !!(size % PAGE_SIZE);
		int err = free_pages(addr, nr_pages);
		if(err) {
			insane(err, "Couldn't free pages: 0x%x, %u pages\n", addr, nr_pages);
		}
		return;
	}

	//Verify that the address is aligned on MEM_BLOCK_SIZE boundary
	ptr_t addr_i = (ptr_t)addr;
	if(!((addr_i % MEM_BLOCK_SIZE) == 0)) {
		insane(ESTINV, "Memory being freed is not block aligned: 0x%x of claimed size %d B\n", addr, size);
	}

	//Get the page address from the addr
	struct mem_page *page = (void *)(addr_i & PAGE_MASK);
	if(page->sign != MEM_PAGE_SIGN) {
		insane(ESTINV, "mem_page signature mismatch for 0x%x of claimed size %d B\n", addr, size);
	}

	//Store the number of free blocks it has before we free
	uint_t init_free_blocks = page->free_blocks;
	//Mark the bit not-free
	uint_t bit = (addr_i - (ptr_t)page) / MEM_BLOCK_SIZE;
	uint_t size_blocks = size / MEM_BLOCK_SIZE + !!(size % MEM_BLOCK_SIZE);

	uint_t i;
	for(i = bit; i < (bit + size_blocks); i++)
		clear_bit(&page->bitmap, i);

	page->free_blocks += size_blocks;
	//See if page is partially/fully free
	if(page->free_blocks == MEM_PAGE_AVL_BLOCKS) {
		//Detach from pool and free it
		list_detach(pages, pages_tail, page);
		free_pages(page, 1);
	} else if(init_free_blocks == 0) {
		//Plug the page in the list of pages if
		//it was out (had no free blocks)
		list_attach(pages, pages_tail, page);
	}
}
